

import numpy as np
import torch
import torch.multiprocessing as mp
import torchaudio
import joblib
import os
import threading
import math
import numpy as np
import itertools
from tqdm import tqdm
from pathlib import Path
import random
import glob
import librosa
import sys
import soundfile as sf
import utils_file
from gxl_wv_mos import get_wvmos
LOGGING_INTERVAL = 5

NUM_THREADS = int(sys.argv[1])
INPUT_SCP = sys.argv[2]  ### filename path\n
OUTPUT_DIR = sys.argv[3]
# RESULT_TEXT = sys.argv[4]


def inference(rank, queue: mp.Queue):
    gpu_id = 0
    out_scp = os.path.join(OUTPUT_DIR, f"wvmos_sub{rank}.scp")
    model = get_wvmos(cuda=True, device=f"cuda:{gpu_id}")
    res = {}
    while True:
        lines = queue.get()
        if lines is None:
            break
        item = lines[0]

        k, v = item
        res[k] = model.calculate_one(v)

    utils_file.write_dic_to_scp(res, out_scp)

def setInterval(interval):
    def decorator(function):
        def wrapper(*args, **kwargs):
            stopped = threading.Event()

            def loop():  # executed in another thread
                while not stopped.wait(interval):  # until stopped
                    function(*args, **kwargs)

            t = threading.Thread(target=loop)
            t.daemon = True  # stop if the program exits
            t.start()
            return stopped

        return wrapper

    return decorator


last_batches = None


@setInterval(LOGGING_INTERVAL)
def QueueWatcher(queue):
    global last_batches
    curr_batches = queue.qsize()
    print(
        f"Remain: {curr_batches} batches [ {(last_batches-curr_batches)/LOGGING_INTERVAL} batches/s ]"
    )
    last_batches = curr_batches


if __name__ == "__main__":
    mp.set_start_method('spawn',force=True)

    print(f"Running with {NUM_THREADS} threads and batchsize 1")
    processes = []
    queue = mp.Queue()
    for rank in range(NUM_THREADS):
        p = mp.Process(target=inference, args=(rank, queue))
        p.start()
        processes.append(p)

    accum = []
    tmp_file = []
    
    dic_data = utils_file.load_dic_from_pathscp(INPUT_SCP)
    for item in tqdm(dic_data.items()):
 
        item = (item[0], item[1].split()[1])
        accum.append(item)
        if len(accum) == 1:
            queue.put(accum.copy())
            accum.clear()


    for _ in range(NUM_THREADS):
        queue.put(None)

    last_batches = queue.qsize()
    queue_watcher = QueueWatcher(queue)
    for p in processes:
        p.join()
    queue_watcher.set()

